#pragma once

#include "NvInferPlugin.h"
#include "plugin.h"
#include <cstdlib>
#include <cublas_v2.h>
#include <cuda_fp16.h>
#include <cudnn.h>
#include <iostream>
#include <string>
#include <thrust/device_ptr.h>
#include <vector>

namespace nvinfer1
{
  namespace plugin
  {

    class Sxai : public nvinfer1::IPluginV2DynamicExt
    {
    public:
      // TODO: more constructor
      Sxai();

      Sxai(const void *data, size_t length);

      ~Sxai() override;

      int getNbOutputs() const noexcept override;

      DimsExprs
      getOutputDimensions(int index, const DimsExprs *inputs, int nbInputDims,
                          nvinfer1::IExprBuilder &exprBuilder) noexcept override;

      int initialize() noexcept override;

      void terminate() noexcept override;

      size_t getWorkspaceSize(const nvinfer1::PluginTensorDesc *inputs,
                              int nbInputs,
                              const nvinfer1::PluginTensorDesc *outputs,
                              int nbOutputs) const noexcept override;

      int enqueue(const nvinfer1::PluginTensorDesc *inputDesc,
                  const nvinfer1::PluginTensorDesc *outputDesc,
                  const void *const *inputs, void *const *outputs, void *workspace,
                  cudaStream_t stream) noexcept override;

      size_t getSerializationSize() const noexcept override;

      void serialize(void *buffer) const noexcept override;

      bool supportsFormatCombination(int pos,
                                     const nvinfer1::PluginTensorDesc *inOut,
                                     int nbInputs, int nbOutputs) noexcept override;

      const char *getPluginType() const noexcept override;

      const char *getPluginVersion() const noexcept override;

      void destroy() noexcept override;

      nvinfer1::IPluginV2DynamicExt *clone() const noexcept override;

      void setPluginNamespace(const char *pluginNamespace) noexcept override;

      const char *getPluginNamespace() const noexcept override;

      DataType getOutputDataType(int index, const nvinfer1::DataType *inputTypes,
                                 int nbInputs) const noexcept override;

      void attachToContext(cudnnContext *cudnn, cublasContext *cublas,
                           nvinfer1::IGpuAllocator *allocator) noexcept override;

      void detachFromContext() noexcept override;

      void configurePlugin(const nvinfer1::DynamicPluginTensorDesc *in,
                           int nbInputs,
                           const nvinfer1::DynamicPluginTensorDesc *out,
                           int nbOutputs) noexcept override;

    protected:
      // attributes load from onnx

      // namespace
      const char *namespace_;
      cublasHandle_t mCublasHandle;
      nvinfer1::DataType mDataType = nvinfer1::DataType::kFLOAT;
    };

    class SxaiPluginCreator : public BaseCreator
    {
    public:
      SxaiPluginCreator();

      ~SxaiPluginCreator() override = default;

      const char *getPluginName() const noexcept override;

      const char *getPluginVersion() const noexcept override;

      const PluginFieldCollection *getFieldNames() noexcept override;

      IPluginV2DynamicExt *
      createPlugin(const char *name,
                   const PluginFieldCollection *fc) noexcept override;

      IPluginV2DynamicExt *deserializePlugin(const char *name,
                                             const void *serialData,
                                             size_t serialLength) noexcept override;

    private:
      static PluginFieldCollection field_collection_;
      static std::vector<PluginField> attributes_;
    };

    REGISTER_TENSORRT_PLUGIN(SxaiPluginCreator);

  } // namespace plugin
} // namespace nvinfer1